import nptdms
import numpy as np
from pathlib import Path
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
from sklearn.neural_network import MLPClassifier
from sklearn.metrics import classification_report

class TDMSNeuralClassifier:
    def __init__(self, hidden_layer_sizes=(100,), max_iter=1000):
        """
        初始化神经网络分类器
        Args:
            hidden_layer_sizes: 隐藏层的神经元数量
            max_iter: 最大迭代次数
        """
        self.scaler = StandardScaler()
        self.classifier = MLPClassifier(
            hidden_layer_sizes=hidden_layer_sizes,
            max_iter=max_iter,
            random_state=42
        )
        self.labels = None
        
    def read_tdms(self, file_path):
        """
        读取TDMS文件并提取特征
        """
        tdms_file = nptdms.TdmsFile.read(file_path)
        features = []
        
        for group in tdms_file.groups():
            for channel in group.channels():
                data = channel[:]
                # 提取基本统计特征
                features.extend([
                    np.mean(data),
                    np.std(data),
                    np.max(data),
                    np.min(data),
                    np.median(data)
                ])
                
        return np.array(features)
    
    def prepare_data(self, data_dir):
        """
        准备训练数据
        """
        features = []
        labels = []
        data_path = Path(data_dir)
        
        # 遍历所有TDMS文件
        for tdms_file in data_path.rglob('*.tdms'):
            try:
                # 获取类别（文件夹名称）
                category = tdms_file.parent.name
                # 读取并提取特征
                feature_vector = self.read_tdms(tdms_file)
                
                features.append(feature_vector)
                labels.append(category)
                
            except Exception as e:
                print(f"处理文件 {tdms_file} 时出错: {e}")
        
        return np.array(features), np.array(labels)
    
    def train(self, data_dir, test_size=0.2):
        """
        训练神经网络
        """
        # 准备数据
        X, y = self.prepare_data(data_dir)
        self.labels = np.unique(y)
        
        # 分割训练集和测试集
        X_train, X_test, y_train, y_test = train_test_split(
            X, y, test_size=test_size, random_state=42
        )
        
        # 标准化特征
        X_train = self.scaler.fit_transform(X_train)
        X_test = self.scaler.transform(X_test)
        
        # 训练模型
        self.classifier.fit(X_train, y_train)
        
        # 评估模型
        y_pred = self.classifier.predict(X_test)
        print("\n分类报告:")
        print(classification_report(y_test, y_pred))
        
        return self.classifier.score(X_test, y_test)
    
    def predict(self, tdms_file_path):
        """
        预测新数据的类别
        """
        # 提取特征
        features = self.read_tdms(tdms_file_path)
        features = features.reshape(1, -1)
        
        # 标准化特征
        features_scaled = self.scaler.transform(features)
        
        # 预测
        prediction = self.classifier.predict(features_scaled)
        probabilities = self.classifier.predict_proba(features_scaled)
        
        # 获取预测概率
        pred_prob = max(probabilities[0])
        
        return prediction[0], pred_prob 